import sys
from Bio import SeqIO
from Bio import Restriction as rst
from Bio.Seq import Seq
import re
from matplotlib import pyplot
import argparse
from collections import Counter
import multiprocessing as mp

parser = argparse.ArgumentParser(description='预测并统计DNA序列内IIB型内切酶位点')
parser.add_argument('-f', '--fasta', type=str, help='DNA序列fasta文件')
parser.add_argument('-e', '--enzyme', type=str, default='BaeI,BcgI,BsaXI,CspCI,AjuI,AloI,BplI,ArsI,BarI,FalI,PsrI', help='逗号分隔的IIB内切酶, '+
'默认BaeI,BcgI,BsaXI,CspCI,AjuI,AloI,BplI,ArsI,BarI,FalI,PsrI; '+
'支持：AjuI,AlfI,AloI,ArsI,BaeI,BarI,BceSIV,BcgI,BdaI,BplI,BsaXI,Bsp24I,CjeI,CjePI,CspCI,FalI,Hin4I,NgoAVIII,NmeDI,PpiI,PsrI,RdeGBIII,SdeOSI,TstI,UcoMSI; '+
'HsoII不在Biopython的Restriction包中支持, HaeIV切割片段的粘性末端不等长, 故不支持; '+
'BaeI BcgI BsaXI CspCI可在New England Biolabs公司购买, '+
'AjuI AloI BplI可在Thermo Fisher Scientific公司购买, '+
'ArsI BarI FalI PsrI可在SibEnzyme Ltd.公司购买')
parser.add_argument('-o', '--out', type=str, default='IIB', help='输出文件前缀, 输出的序列文件中, 序列名末尾为‘_F’为内切酶识别位点在序列正向, 序列名末尾为‘_R’为内切酶识别位点在序列反向, 序列名末尾为‘_G’为序列长度与内切酶识别位点等长而不具有内切酶识别位点的序列')
parser.add_argument('-p', '--process', type=int, default=11, help='线程数一般不多于内切酶数量, 默认11')
parser.add_argument('-d', '--distance', action='store_true', help='是否输出酶切位点间距长度的堆积图')
args = parser.parse_args()
inputfile = args.fasta
outputfile = args.out
e = args.enzyme
en = e.split(',')
p = args.process
distance = args.distance

IIB_dic = {'AjuI':[(7,12), 'GAANNNNNNNTTGG', (11,6)],
'AlfI':[(10,12), 'GCANNNNNNTGC', (12,10)],
'AloI':[(7,12), 'GAACNNNNNNTCC', (12,7)],
'ArsI':[(8,13), 'GACNNNNNNTTYG', (11,6)],
'BaeI':[(10,15), 'ACNNNNGTAYC', (12,7)],
'BarI':[(7,12), 'GAAGNNNNNNTAC', (12,7)],
'BceSIV':[(7,5), 'GCAGC', (9,11)],
'BcgI':[(10,12), 'CGANNNNNNTGC', (12,10)],
'BdaI':[(10,12), 'TGANNNNNNTCA', (12,10)],
'BplI':[(8,13), 'GAGNNNNNCTC', (13,8)],
'BsaXI':[(9,12), 'ACNNNNNCTCC', (10,7)],
'Bsp24I':[(8,13), 'GACNNNNNNTGG', (12,7)],
'CjeI':[(8,14), 'CCANNNNNNGT', (15,9)],
'CjePI':[(7,13), 'CCANNNNNNNTC', (14,8)],
'CspCI':[(11,13), 'CAANNNNNGTGG', (12,10)],
'FalI':[(8,13), 'AAGNNNNNCTT', (13,8)],
'Hin4I':[(8,13), 'GAYNNNNNVTC', (13,8)],
'NgoAVIII':[(12,14), 'GACNNNNNTGA', (13,11)],
'NmeDI':[(12,7), 'RCCGGY', (7,12)],
'PpiI':[(7,12), 'GAACNNNNNCTC', (13,8)],
'PsrI':[(7,12), 'GAACNNNNNNTAC', (12,7)],
'RdeGBIII':[(9,11), 'TGRYCA', (11,9)],
'SdeOSI':[(11,13), 'GACNNNNRTGA', (12,10)],
'TstI':[(8,13), 'CACNNNNNNTCC', (12,7)],
'UcoMSI':[(7,5), 'GAGCTC', (5,7)]}

for e in en:
    if not e in IIB_dic:
        print('输入的序列不是IIB内切酶')
        sys.exit(0)

def IIB_perdict(i,o,e,d):
    out_2b = open(o+'.'+str(e)+'.fas', "w") # 输出2b内切酶预测序列
    dist = open(o+'.'+str(e)+'.distance.txt', "w") # 输出间隔长度
    stat = open(o+'.'+str(e)+'.stat.txt', "w") # 输出统计结果
    rb = rst.RestrictionBatch([e])
    en_dif = abs(IIB_dic[e][0][0]-IIB_dic[e][0][1]) # 粘性末端长度
    en_len = len(rb.get(e).site) + int((IIB_dic[e][0][0]+IIB_dic[e][0][1]+IIB_dic[e][2][0]+IIB_dic[e][2][1])/2) + en_dif # 切割总长度
    stat.write('Cohesive end length='+str(en_dif)+'\tFragment length='+str(en_len)+'\n')
    print(e,(''.join([str(IIB_dic[e][0]),IIB_dic[e][1],str(IIB_dic[e][2])])).replace(',','/',2)) # 打印内切酶
    
    pattern_f = rb.get(e).site # 生成正向序列正则模式
    pattern_f = re.sub('R','[AG]',pattern_f)
    pattern_f = re.sub('Y','[CT]',pattern_f)
    pattern_f = re.sub('M','[AC]',pattern_f)
    pattern_f = re.sub('K','[GT]',pattern_f)
    pattern_f = re.sub('S','[GC]',pattern_f)
    pattern_f = re.sub('W','[AT]',pattern_f)
    pattern_f = re.sub('H','[ATC]',pattern_f)
    pattern_f = re.sub('B','[GTC]',pattern_f)
    pattern_f = re.sub('V','[GAC]',pattern_f)
    pattern_f = re.sub('D','[GAT]',pattern_f)
    pattern_f = re.sub('N','[ATCG]',pattern_f)
    pattern_f = '[A-Z]*'+pattern_f
    
    pattern_r = rb.get(e).site # 序列反向
    pattern_r = Seq(pattern_r)
    pattern_r = pattern_r.reverse_complement()
    pattern_r = str(pattern_r)
    
    pattern_r = re.sub('R','[AG]',pattern_r) # 生成反向序列正则模式
    pattern_r = re.sub('Y','[CT]',pattern_r)
    pattern_r = re.sub('M','[AC]',pattern_r)
    pattern_r = re.sub('K','[GT]',pattern_r)
    pattern_r = re.sub('S','[GC]',pattern_r)
    pattern_r = re.sub('W','[AT]',pattern_r)
    pattern_r = re.sub('H','[ATC]',pattern_r)
    pattern_r = re.sub('B','[GTC]',pattern_r)
    pattern_r = re.sub('V','[GAC]',pattern_r)
    pattern_r = re.sub('D','[GAT]',pattern_r)
    pattern_r = re.sub('N','[ATCG]',pattern_r)
    pattern_r = '[A-Z]*'+pattern_r
    
    cutlist = [] # 所有序列存入列表
    distancelist = [] # 所有间距存入列表
    for sequence in SeqIO.parse(i, 'fasta'):
        seq_id = sequence.id
        seq2cut = Seq(str(sequence.seq))
        allsite = []
        for i in rb.search(seq2cut):
            allsite += rb.search(seq2cut)[i]
        allsite.sort
        for n in range(0,len(allsite)-2): # 
            start = allsite[n]-en_dif
            end = allsite[n+1]-1
            if end - start + 1 == en_len:
                cutseq = sequence[start-1:end].upper()
                if	re.match(pattern_f,str(cutseq.seq)):
                    print('>'+seq_id+"_"+str(start)+"_"+str(end)+"_F", cutseq.seq, file = out_2b, sep = '\n')
                    cutlist.append(str(cutseq.seq))
                elif re.match(pattern_r,str(cutseq.seq)):
                    cutseq = cutseq.reverse_complement()
                    print('>'+seq_id+"_"+str(end)+"_"+str(start)+"_R", cutseq.seq, file = out_2b, sep = '\n')
                    cutlist.append(str(cutseq.seq))
                else:
                    print('>'+seq_id+"_"+str(start)+"_"+str(end)+"_G", cutseq.seq, file = out_2b, sep = '\n') #片段长度恰好等于酶切片段长度而不含有IIB内切酶识别序列
                    cutlist.append(str(cutseq.seq))
            else:
                distancelist.append(allsite[n+1]-en_dif-allsite[n])
                dist.write(str(allsite[n+1]-en_dif-allsite[n])+"\t")
    
    all_seq_num = len(cutlist)
    unique_seq_num = len(set(cutlist))
    stat.write('All seq number: '+str(all_seq_num)+'\tAvailable seq number: '+str(unique_seq_num)+'\n')
    
    seqstat = dict(Counter(cutlist))
    s = 0
    m = 0
    for i in seqstat:			
        if seqstat[i] != 1:
            m += 1
            stat.write(str(seqstat[i])+'\t'+str(i)+'\n')
        else:
            s += 1
    stat.write('Unique site:'+str(s)+'\tMaltiple sequence:'+str(m))
    
    if d == True:
        pyplot.hist(distancelist, bins=100, color='blue', histtype='stepfilled')
        pyplot.xlabel('Distance Length',fontname="Times New Roman",size=18)
        pyplot.ylabel('Distance Frequency',fontname="Times New Roman",size=18)
        pyplot.title('Histogram of '+e+' Distance',fontname="Times New Roman",size=18)
        pyplot.show()

def job_map(l):
    return IIB_perdict(l[0], l[1], l[2], l[3])

if __name__ == '__main__':
    parameter = [(inputfile,outputfile,i,distance) for i in en]
    if p > mp.cpu_count():
        p = mp.cpu_count()
    with mp.Pool(processes=p) as pool:
        result = pool.map(job_map, parameter)
        pool.close()